#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
固件升级器模块

该模块提供以下功能:
1. 处理固件升级的核心逻辑
2. 发送握手包并等待设备响应
3. 发送BIN文件数据包
4. 处理重试机制和超时控制
"""

import json
import time
import socket
import tkinter as tk
from tkinter import ttk, messagebox
from typing import Optional, Callable
import queue

# 导入现有的JSON发送器和接收器
from json_bin_sender import JSONBinSender, JSONBinReceiver


class FirmwareUpgrader:
    """固件升级器类"""
    
    def __init__(self, socket_connection, bin_loader):
        """
        初始化固件升级器
        
        Args:
            socket_connection: 已建立的socket连接
            bin_loader: BIN加载器实例
        """
        self.socket = socket_connection
        self.bin_loader = bin_loader
        
        # 超时配置（默认值，可根据需要调整）
        self.CONNECT_TIMEOUT = 30.0      # 连接超时时间（秒）
        self.HANDSHAKE_TIMEOUT = 30.0    # 握手超时时间（秒）
        self.DATA_SEND_TIMEOUT = 600.0    # 数据发送超时时间（秒）
        self.RESPONSE_TIMEOUT = 30.0     # 等待设备响应超时时间（秒）
        
        # 回调函数
        self.progress_callback = None
        self.message_callback = None
        self.upgrade_status_callback = None
        
        # 升级状态
        self.upgrade_in_progress = False

        self._message_queue = queue.Queue()  # 新增：消息队列
    
    def set_timeouts(self, connect_timeout=30.0, handshake_timeout=30.0, 
                     data_send_timeout=30.0, response_timeout=30.0):
        """
        设置超时参数
        
        Args:
            connect_timeout (float): 连接超时时间
            handshake_timeout (float): 握手超时时间
            data_send_timeout (float): 数据发送超时时间
            response_timeout (float): 响应超时时间
        """
        self.CONNECT_TIMEOUT = connect_timeout
        self.HANDSHAKE_TIMEOUT = handshake_timeout
        self.DATA_SEND_TIMEOUT = data_send_timeout
        self.RESPONSE_TIMEOUT = response_timeout
    
    def set_callbacks(self, progress_callback=None, message_callback=None, 
                     upgrade_status_callback=None):
        """
        设置回调函数
        
        Args:
            progress_callback: 进度更新回调函数
            message_callback: 消息显示回调函数
            upgrade_status_callback: 升级状态回调函数
        """
        self.progress_callback = progress_callback
        self.message_callback = message_callback
        self.upgrade_status_callback = upgrade_status_callback
    
    def on_receive_message(self, message):
        """由外部调用，注入收到的回包消息"""
        self._message_queue.put(message)

    def _wait_for_response(self, timeout=5.0):
        """等待设备响应，优先从队列取消息"""
        try:
            message = self._message_queue.get(timeout=timeout)
            return message
        except queue.Empty:
            return None

    def send_handshake(self):
        """
        发送握手包并等待回包
        
        Returns:
            bool: 握手成功返回True，否则返回False
        """
        try:
            # 发送JSON格式的握手包
            handshake_packet = {"cmd": 996, "script": 9}
            handshake_json = json.dumps(handshake_packet)
            self.socket.sendall(handshake_json.encode('utf-8'))
            self._display_message(f"已发送握手包: {handshake_json}，等待设备响应...")

            response_json = self._wait_for_response(self.HANDSHAKE_TIMEOUT)
            if response_json is None:
                self._display_message("握手超时，设备未响应")
                return False

            # 使用JSON接收器解析回包
            json_receiver = JSONBinReceiver()
            try:
                response = json.loads(response_json)
                # 首先检查是否是cmd996响应
                if int(response.get("cmd", 0)) != 996:
                    self._display_message(f"忽略非cmd996响应: {response_json}")
                    return False
                
                # 使用JSON接收器严格解析
                success, result = json_receiver.parse_response(response_json)
                if success and result.get("result") in [1, 2]:
                    self._display_message("握手成功，设备已准备就绪")
                    return True
                else:
                    self._display_message(f"握手失败，cmd996响应无效: {result.get('error', '未知错误')}")
                    return False
            except json.JSONDecodeError:
                self._display_message(f"无效的JSON响应: {response_json}")
                return False
            except Exception as e:
                self._display_message(f"解析响应时出错: {str(e)}")
                return False
        except Exception as e:
            self._display_message(f"握手过程中发生错误: {str(e)}")
            return False

    def send_bin_data(self, start_address: int = 0x08000000, block_length: int = 256):
        """
        发送BIN文件数据，使用JSON格式，每次8字节
        """
        try:
            # 获取BIN数据
            if not hasattr(self.bin_loader, 'bin_loader'):
                self._display_message("BIN加载器无效")
                return False
                
            try:
                # 调试信息：检查bin_loader状态
                self._display_message(f"BIN加载器状态: {type(self.bin_loader.bin_loader)}")
                if hasattr(self.bin_loader.bin_loader, 'file_path'):
                    self._display_message(f"已加载BIN文件: {self.bin_loader.bin_loader.file_path}")
                else:
                    self._display_message("未检测到已加载的BIN文件")
                
                bin_data = self.bin_loader.bin_loader.get_memory_data()
                if not bin_data:
                    self._display_message("BIN数据为空")
                    return False
                self._display_message(f"成功加载BIN数据，大小: {len(bin_data)}字节")
            except Exception as e:
                self._display_message(f"加载BIN数据失败: {str(e)}")
                return False
            
            total_bytes = len(bin_data)
            self._display_message(f"开始发送BIN数据，总大小: {total_bytes} 字节")
            
            # 创建JSON BIN发送器
            chip_type = getattr(self, "chip_type", "ch32v10x")
            json_sender = JSONBinSender(bin_data, start_address, block_length, chip_type)
            json_receiver = JSONBinReceiver()
            
            # 发送数据包
            sent_bytes = 0
            start_time = time.time()
            while sent_bytes < total_bytes:
                if not self.upgrade_in_progress:
                    return False
                
                packet = json_sender.get_next_packet()
                if packet is None:
                    break
                
                packet_json = json.dumps(packet)
                try:
                    self.socket.sendall(packet_json.encode('utf-8'))
                    self._display_message(f"成功发送数据包: {packet_json[:100]}...")  # 只显示前100字符
                except Exception as e:
                    self._display_message(f"发送数据包失败: {str(e)}")
                    return False
                
                response_received = False
                retry_count = 0
                max_retries = 3
                while not response_received and self.upgrade_in_progress and retry_count < max_retries:
                    response_json = self._wait_for_response(5.0)
                    if response_json is None:
                        retry_count += 1
                        if retry_count < max_retries:
                            self._display_message(f"等待设备响应超时 (第{retry_count}次重试)")
                            self.socket.sendall(packet_json.encode('utf-8'))
                            self._display_message(f"重新发送数据包: {packet_json}")
                        else:
                            self._display_message("设备响应超时，超过最大重试次数")
                            return False
                        continue

                    try:
                        response = json.loads(response_json)
                        # 忽略非cmd996响应
                        if int(response.get("cmd", 0)) != 996:
                            self._display_message(f"收到非目标指令，继续等待: {response_json}")
                            continue
                            
                        success, result = json_receiver.parse_response(response_json)
                        if success and result.get("result") in [1, 2]:
                            if result.get("script") == packet.get("script"):
                                response_received = True
                                self._display_message(f"收到有效响应: {response_json}")
                            else:
                                self._display_message(
                                    f"设备响应script字段不匹配，期望 {packet.get('script')}, 实际 {result.get('script')}，继续等待...")
                                continue
                        else:
                            self._display_message(f"设备响应错误: {result.get('error', '未知错误')}")
                            return False
                    except json.JSONDecodeError:
                        self._display_message(f"无效的JSON响应: {response_json}")
                        continue
                    except Exception as e:
                        self._display_message(f"解析响应时出错: {str(e)}")
                        continue

                # 更新进度（每次发送8字节）
                sent_bytes += 8
                self._update_progress(min(sent_bytes, total_bytes), total_bytes)
            
            # 发送确认包
            confirm_packet = {"cmd": 996, "script": 1, "result": 2}
            confirm_json = json.dumps(confirm_packet)
            self.socket.sendall(confirm_json.encode('utf-8'))
            self._display_message("已发送数据确认包 (script=1)")
            
            # 等待确认响应
            response_json = self._wait_for_response(5.0)
            if response_json is None:
                self._display_message("确认响应超时")
                return False
            
            # 发送复位包
            reset_packet = {"cmd": 996, "script": 3}
            reset_json = json.dumps(reset_packet)
            self.socket.sendall(reset_json.encode('utf-8'))
            self._display_message("已发送设备复位指令 (script=3)，等待设备重启升级")
            
            return True
            
        except Exception as e:
            self._display_message(f"发送BIN数据时发生错误: {str(e)}")
            return False
    
    def execute_upgrade_process(self, start_address: int = 0x08000000, block_length: int = 256):
        """
        执行升级过程
        
        Args:
            start_address (int): 起始地址
            
        Returns:
            bool: 升级成功返回True，否则返回False
        """
        try:
            self.upgrade_in_progress = True
            self._update_upgrade_status(True)
            
            # 显示升级开始信息
            self._display_message("开始固件升级...")
            
            # 1. 发送握手包并等待回包
            if not self.send_handshake():
                self._upgrade_failed("握手失败")
                return False
            
            # 2. 发送BIN文件数据
            if not self.send_bin_data(start_address, block_length):
                self._upgrade_failed("发送BIN数据失败")
                return False
            
            # 升级成功
            self._upgrade_success()
            return True
            
        except Exception as e:
            self._upgrade_failed(f"升级过程中发生错误: {str(e)}")
            return False
        finally:
            self.upgrade_in_progress = False
            self._update_upgrade_status(False)
    
    def cancel_upgrade(self):
        """取消升级"""
        self.upgrade_in_progress = False
    
    def is_upgrading(self):
        """
        检查是否正在升级
        
        Returns:
            bool: 正在升级返回True，否则返回False
        """
        return self.upgrade_in_progress
    
    def _update_progress(self, sent_bytes, total_bytes):
        """
        更新进度
        
        Args:
            sent_bytes (int): 已发送字节数
            total_bytes (int): 总字节数
        """
        if self.progress_callback:
            self.progress_callback(sent_bytes, total_bytes)
    
    def _display_message(self, message):
        """
        显示消息
        
        Args:
            message (str): 要显示的消息
        """
        if self.message_callback:
            self.message_callback(message)
    
    def _update_upgrade_status(self, in_progress):
        """
        更新升级状态
        
        Args:
            in_progress (bool): 是否正在升级
        """
        if self.upgrade_status_callback:
            self.upgrade_status_callback(in_progress)
    
    def _upgrade_success(self):
        """升级成功处理"""
        self._display_message("固件升级成功完成")
    
    def _upgrade_failed(self, error_message):
        """升级失败处理"""
        self._display_message(f"固件升级失败: {error_message}")


# 使用示例和测试函数
def main():
    """主函数，用于测试和演示"""
    print("固件升级器模块演示")
    print("该模块需要与GUI应用程序集成使用")

if __name__ == "__main__":
    main()
